{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import pickle\n",
    "import numpy as np\n",
    "import os\n",
    "import json\n",
    "import random\n",
    "    \n",
    "# load source words\n",
    "source_words_path = os.path.join(os.getcwd(), 'source_words.pkl')\n",
    "with open(source_words_path, 'rb') as f_source_words:\n",
    "    source_words = pickle.load(f_source_words)\n",
    "    \n",
    "# load target words\n",
    "target_words_path = os.path.join(os.getcwd(), 'target_words.pkl')\n",
    "with open(target_words_path, 'rb') as f_target_words:\n",
    "    target_words = pickle.load(f_target_words)\n",
    "    \n",
    "# load label words\n",
    "label_words_path = os.path.join(os.getcwd(), 'label_words.pkl')\n",
    "with open(label_words_path, 'rb') as f_label_words:\n",
    "    label_words = pickle.load(f_label_words)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "945\n",
      "133\n",
      "27\n",
      "1\n",
      "3\n",
      "2\n",
      "0\n",
      "1\n"
     ]
    }
   ],
   "source": [
    "print(len(source_words))\n",
    "print(len(target_words))\n",
    "print(len(label_words))\n",
    "print(source_words['<pad>'])\n",
    "print(source_words['<eos>'])\n",
    "print(source_words['<sos>'])\n",
    "print(source_words['<unk>'])\n",
    "print(target_words['<pad>'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "编码器Encoder的实现\n",
    "'''\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, intent_dim, hid_dim, n_layers, kernel_size, dropout, max_length=50):\n",
    "        super(Encoder, self).__init__()\n",
    "        \n",
    "        assert kernel_size % 2 == 1,'kernel size must be odd!' # 卷积核size为奇数，方便序列两边pad处理\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device) # 确保整个网络的方差不会发生显著变化\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(input_dim, emb_dim) # token编码\n",
    "        self.pos_embedding = nn.Embedding(max_length, emb_dim) # token的位置编码\n",
    "        \n",
    "        self.emb2hid = nn.Linear(emb_dim, hid_dim) # 线性层，从emb_dim转为hid_dim\n",
    "        self.hid2emb = nn.Linear(hid_dim, emb_dim) # 线性层，从hid_dim转为emb_dim\n",
    "        \n",
    "        # 卷积块\n",
    "        self.convs = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                              out_channels=2*hid_dim, # 卷积后输出的维度，这里2*hid_dim是为了后面的glu激活函数\n",
    "                                              kernel_size=kernel_size,\n",
    "                                              padding=(kernel_size - 1)//2) # 序列两边补0个数，保持维度不变\n",
    "                                              for _ in range(n_layers)]) \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        # 利用encoder的输出进行意图识别\n",
    "        self.intent_output = nn.Linear(emb_dim, intent_dim)\n",
    "        \n",
    "    def forward(self, src):\n",
    "        # src: [batch_size, src_len]\n",
    "        batch_size = src.shape[0]\n",
    "        src_len = src.shape[1]\n",
    "        \n",
    "        # 创建token位置信息\n",
    "        pos = torch.arange(src_len).unsqueeze(0).repeat(batch_size, 1).to(device) # [batch_size, src_len]\n",
    "        \n",
    "        # 对token与其位置进行编码\n",
    "        tok_embedded = self.tok_embedding(src) # [batch_size, src_len, emb_dim]\n",
    "        pos_embedded = self.pos_embedding(pos.long()) # [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        # 对token embedded和pos_embedded逐元素加和\n",
    "        embedded = self.dropout(tok_embedded + pos_embedded) # [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        # embedded经过一线性层，将emb_dim转为hid_dim，作为卷积块的输入\n",
    "        conv_input = self.emb2hid(embedded) # [batch_size, src_len, hid_dim]\n",
    "        \n",
    "        # 转变维度，卷积在输入数据的最后一维进行\n",
    "        conv_input = conv_input.permute(0, 2, 1) # [batch_size, hid_dim, src_len]\n",
    "        \n",
    "        # 以下进行卷积块\n",
    "        for i, conv in enumerate(self.convs):\n",
    "            # 进行卷积\n",
    "            conved = conv(self.dropout(conv_input)) # [batch_size, 2*hid_dim, src_len]\n",
    "            \n",
    "            # 进行激活glu\n",
    "            conved = F.glu(conved, dim=1) # [batch_size, hid_dim, src_len]\n",
    "            \n",
    "            # 进行残差连接\n",
    "            conved = (conved + conv_input) * self.scale # [batch_size, hid_dim, src_len]\n",
    "            \n",
    "            # 作为下一个卷积块的输入\n",
    "            conv_input = conved\n",
    "        \n",
    "        # 经过一线性层，将hid_dim转为emb_dim，作为enocder的卷积输出的特征\n",
    "        conved = self.hid2emb(conved.permute(0, 2, 1)) # [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        # 又是一个残差连接，逐元素加和输出，作为encoder的联合输出特征\n",
    "        combined = (conved + embedded) * self.scale # [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        # 意图识别,加一个平均池化,池化后的维度是：[batch_size, emb_dim]\n",
    "        intent_output = self.intent_output(F.avg_pool1d(combined.permute(0, 2, 1), combined.shape[1]).squeeze()) # [batch_size, intent_dim]\n",
    "        \n",
    "        return conved, combined, intent_output\n",
    "    \n",
    "'''\n",
    "解码器Decoder实现\n",
    "'''\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, emb_dim, hid_dim, n_layers,kernel_size, dropout, trg_pad_idx, max_length=50):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.kernel_size = kernel_size\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([0.5])).to(device)\n",
    "        \n",
    "        self.tok_embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        self.pos_embedding = nn.Embedding(max_length, emb_dim)\n",
    "        \n",
    "        self.emb2hid = nn.Linear(emb_dim, hid_dim)\n",
    "        self.hid2emb = nn.Linear(hid_dim, emb_dim)\n",
    "        \n",
    "        self.attn_hid2emb = nn.Linear(hid_dim, emb_dim)\n",
    "        self.attn_emb2hid = nn.Linear(emb_dim, hid_dim)\n",
    "        \n",
    "        self.fc_out = nn.Linear(emb_dim, output_dim)\n",
    "        \n",
    "        self.convs = nn.ModuleList([nn.Conv1d(in_channels=hid_dim,\n",
    "                                              out_channels=2*hid_dim,\n",
    "                                              kernel_size=kernel_size)\n",
    "                                              for _ in range(n_layers)])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def calculate_attention(self, embedded, conved, encoder_conved, encoder_combined):\n",
    "        '''\n",
    "        embedded:[batch_size, trg_Len, emb_dim]\n",
    "        conved:[batch_size, hid_dim, trg_len]\n",
    "        encoder_conved:[batch_size, src_len, emb_dim]\n",
    "        encoder_combined:[batch_size, src_len, emb_dim]\n",
    "        '''\n",
    "        # 经过一线性层，将hid_dim转为emb_dim，作为deocder的卷积输出的特征\n",
    "        conved_emb = self.attn_hid2emb(conved.permute(0, 2, 1)) # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # 一个残差连接，逐元素加和输出，作为decoder的联合输出特征\n",
    "        combined = (conved_emb + embedded) * self.scale # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # decoder的联合特征combined与encoder的卷积输出进行矩阵相乘\n",
    "        energy = torch.matmul(combined, encoder_conved.permute(0, 2, 1)) # [batch_size, trg_len, src_len]\n",
    "        \n",
    "        attention = F.softmax(energy, dim=2) # [batch_size, trg_len, src_len]\n",
    "        \n",
    "        attention_encoding = torch.matmul(attention, encoder_combined) # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # 经过一线性层，将emb_dim转为hid_dim\n",
    "        attended_encoding = self.attn_emb2hid(attention_encoding) # [batch_size, trg_len, hid_dim]\n",
    "        \n",
    "        # 一个残差连接，逐元素加和输出\n",
    "        attended_combined = (conved + attended_encoding.permute(0, 2, 1)) * self.scale # [batch_size, hid_dim, trg_len]\n",
    "        \n",
    "        return attention, attended_combined\n",
    "    \n",
    "    def forward(self, trg, encoder_conved, encoder_combined):\n",
    "        '''\n",
    "        trg:[batch_size, trg_len]\n",
    "        encoder_conved:[batch_size, src_len, emb_dim]\n",
    "        encoder_combined:[batch_size, src_len, emb_dim]\n",
    "        '''\n",
    "        batch_size = trg.shape[0]\n",
    "        trg_len = trg.shape[1]\n",
    "        \n",
    "        # 位置编码\n",
    "        pos = torch.arange(trg_len).unsqueeze(0).repeat(batch_size, 1).to(device) # [batch_size, trg_len]\n",
    "        \n",
    "        # 对token和pos进行embedding\n",
    "        tok_embedded = self.tok_embedding(trg) # [batch_size, trg_len, emb_dim]\n",
    "        pos_embedded = self.pos_embedding(pos.long()) # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # 对token embedded和pos_embedded逐元素加和\n",
    "        embedded = self.dropout(tok_embedded + pos_embedded) # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # 经过一线性层，将emb_dim转为hid_dim，作为卷积的输入\n",
    "        conv_input = self.emb2hid(embedded) # [batch_size, trg_len, hid_dim]\n",
    "        \n",
    "        # 转变维度，卷积在输入数据的最后一维进行\n",
    "        conv_input = conv_input.permute(0, 2, 1) # [batch_size, hid_dim, trg_len]\n",
    "        \n",
    "        batch_size = conv_input.shape[0]\n",
    "        hid_dim = conv_input.shape[1]\n",
    "        \n",
    "        # 卷积块\n",
    "        for i, conv in enumerate(self.convs):\n",
    "            conv_input = self.dropout(conv_input)\n",
    "            \n",
    "            # 在序列的一端进行pad\n",
    "            padding = torch.zeros(batch_size, hid_dim, self.kernel_size - 1).fill_(self.trg_pad_idx).to(device)\n",
    "            \n",
    "            padded_conv_input = torch.cat((padding, conv_input), dim=2) # [batch_size, hid_dim, trg_len + kernel_size - 1]\n",
    "            \n",
    "            # 进行卷积\n",
    "            conved = conv(padded_conv_input) # [batch_size, 2 * hid_dim, trg_len]\n",
    "            \n",
    "            # 经过glu激活\n",
    "            conved = F.glu(conved, dim=1) # [batch_size, hid_dim, trg_len]\n",
    "            \n",
    "            # 计算attention\n",
    "            attention, conved = self.calculate_attention(embedded, conved, encoder_conved, encoder_combined) # [batch_size, trg_len, src_len], [batch_size, hid_dim, trg_len]\n",
    "            \n",
    "            # 残差连接\n",
    "            conved = (conved + conv_input) * self.scale # [batch_size, hid_dim, trg_len]\n",
    "            \n",
    "            # 作为下一层卷积的输入\n",
    "            conv_input = conved\n",
    "        \n",
    "        conved = self.hid2emb(conved.permute(0, 2, 1)) # [batch_size, trg_len, emb_dim]\n",
    "        \n",
    "        # 预测输出\n",
    "        output = self.fc_out(self.dropout(conved)) # [batch_size, trg_len, output_dim]\n",
    "        \n",
    "        return output, attention\n",
    "            \n",
    "# 包装Encoder与Decoer\n",
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder):\n",
    "        super(Seq2Seq, self).__init__()\n",
    "        \n",
    "        # 编码器\n",
    "        self.encoder = encoder\n",
    "        \n",
    "        # 解码器用于slot槽识别\n",
    "        self.decoder = decoder\n",
    "        \n",
    "    def forward(self, src, trg):\n",
    "        '''\n",
    "        src:[batch_size, src_len]\n",
    "        trg:[batch_size, trg_Len-1] # decoder的输入去除了<eos>\n",
    "        \n",
    "        encoder_conved是encoder中最后一个卷积层的输出\n",
    "        encoder_combined是encoder_conved + (src_embedding + postional_embedding)\n",
    "        '''\n",
    "        encoder_conved, encoder_combined, intent_output = self.encoder(src) # [batch_size, src_len, emb_dim]; [batch_size, src_len, emb_dim]\n",
    "        \n",
    "        # decoder是对一批数据进行预测输出\n",
    "        slot_output, attention = self.decoder(trg, encoder_conved, encoder_combined) # [batch_size, trg_len-1, output_dim]; [batch_size, trg_len-1, src_len]\n",
    "        \n",
    "        return intent_output, slot_output, attention"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "emb_dim = 64\n",
    "hid_dim = 32\n",
    "enc_layers = 5\n",
    "dec_layers = 5\n",
    "enc_kernel_size = 3\n",
    "dec_kernel_size = 3\n",
    "enc_dropout = 0.25\n",
    "dec_dropout = 0.25\n",
    "\n",
    "model_path = os.path.join(os.getcwd(), \"model.h5\")\n",
    "\n",
    "input_dim = len(source_words) # source 词典大小（即词数量）\n",
    "output_dim = len(target_words) # target 词典大小（即实体类型数量）\n",
    "intent_dim = len(label_words) # label 词典大小（即意图类别数量）\n",
    "\n",
    "trg_pad_idx = target_words['<pad>']\n",
    "enc = Encoder(input_dim, emb_dim, intent_dim, hid_dim, enc_layers, enc_kernel_size, enc_dropout)\n",
    "dec = Decoder(output_dim, emb_dim, hid_dim, dec_layers, dec_kernel_size, dec_dropout, trg_pad_idx)\n",
    "\n",
    "model = Seq2Seq(enc, dec).to(device)\n",
    "\n",
    "model.load_state_dict(torch.load(model_path))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def predict_slot_intent(sentence, model):\n",
    "    \n",
    "    model.eval()\n",
    "\n",
    "    tokenized = sentence.split()  # tokenize the sentence\n",
    "    tokenized = ['<sos>'] + tokenized + ['<eos>']\n",
    "    indexed = [source_words[t] for t in tokenized]  # convert to integer sequence\n",
    "\n",
    "    print(tokenized)\n",
    "    print(indexed)\n",
    "    \n",
    "    src_tensor = torch.LongTensor(indexed)  # convert to tensor\n",
    "    src_tensor = src_tensor.unsqueeze(0).to(device)  # reshape in form of batch,no. of words\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        encoder_conved, encoder_combined, intent_output = model.encoder(src_tensor)\n",
    "    intent_output = intent_output.squeeze()\n",
    "    intent_output = intent_output.argmax()\n",
    "    intent = intent_output.detach().item()\n",
    "    \n",
    "    intent_label = label_words.itos[intent]\n",
    "    \n",
    "    trg_indexes = [target_words['<sos>']]\n",
    "    print('seqence length: {}'.format(src_tensor.shape[1]))\n",
    "    for i in range(1, src_tensor.shape[1]):\n",
    "        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)\n",
    "        \n",
    "        with torch.no_grad():\n",
    "            slot_output, attention = model.decoder(trg_tensor, encoder_conved, encoder_combined)\n",
    "            \n",
    "        pred_token = slot_output.argmax(2)[:, -1].item()\n",
    "        \n",
    "        trg_indexes.append(pred_token)\n",
    "    \n",
    "    trg_tokens = [target_words.itos[i] for i in trg_indexes]\n",
    "    \n",
    "    print('slot prediction: {}'.format(trg_tokens[1:]))\n",
    "    print('intent prediction: {}'.format(intent_label))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['<sos>', 'i', 'would', 'like', 'to', 'find', 'a', 'flight', 'from', 'charlotte', 'to', 'las', 'vegas', 'that', 'makes', 'a', 'stop', 'in', 'st.', 'louis', '<eos>']\n",
      "[2, 13, 40, 29, 4, 87, 16, 11, 5, 100, 4, 90, 89, 34, 345, 16, 127, 18, 67, 144, 3]\n",
      "seqence length: 21\n",
      "slot prediction: ['o', 'o', 'o', 'o', 'o', 'o', 'o', 'o', 'b-fromloc.city_name', 'o', 'b-toloc.city_name', 'i-toloc.city_name', 'o', 'o', 'o', 'o', 'o', 'b-stoploc.city_name', 'i-stoploc.city_name', 'o']\n",
      "intent prediction: flight\n"
     ]
    }
   ],
   "source": [
    "sentence = \"i would like to find a flight from charlotte to las vegas that makes a stop in st. louis\"\n",
    "sentence2 = \"which airlines have first class flights today\"\n",
    "\n",
    "predict_slot_intent(sentence, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
